#! /usr/bin/env python
# coding=utf-8



import os
from Bio import SeqIO

import sys
import argparse



parser = argparse.ArgumentParser(
    description=''' 对过滤后的 Alicut 路径 建立超矩阵
    用法:
    build_super_matrix.py -i Orthogroup_Sequences_mafft_raw_filtered -o Orthogroup_Sequences_mafft_raw_filtered.fa -p Orthogroup_Sequences_mafft_raw_filtered.nex -t charset
    由大天才于2021年7月20日创建于浙江农业大学''')



parser.add_argument('-i',
                help='必须给定，过滤之后ALICUT的文件路径')

parser.add_argument('-o',
                help='必须给定，输出的fasta文件')

parser.add_argument('-p',
                help='必须给定，输出的partition文件')


parser.add_argument('-t',
                help='必须给定，输出的partition的类型')


args = parser.parse_args()

if not args.i or not args.o or not args.p or not args.t:
    parser.print_help()
    sys.exit()

infile = args.i


outfile = args.o

outfile_partition = args.p

tp = args.t

# infile = r'Orthogroup_Sequences_mafft_raw_filtered'

# outfile = r'Orthogroup_Sequences_mafft_raw_filtered.fa'

# outfile_partition = r'Orthogroup_Sequences_mafft_raw_filtered.pa'

# 建立字典

spe_lista = []

partition_lista = []

tmp_dic = {}


for i in os.listdir(infile):

	tmp_dic[i] = {}

	for j in SeqIO.parse(infile+'/'+i,'fasta'):

		spe = str(j.name).split('@')[0]
		
		tmp_dic[i][spe] = str(j.seq)
		
		if spe not in spe_lista:

			spe_lista.append(spe)

# 建立超举证

super_matrix_dic = {spe:'' for spe in spe_lista}

n = 1

start = 1
for i in tmp_dic:

	l = 0

	for j in tmp_dic[i]:

		l = len(tmp_dic[i][j])

	for j in spe_lista:

		if j not in tmp_dic[i]:

			seq = '-'*l
		else:
			seq = tmp_dic[i][j]

		super_matrix_dic[j] += seq

	partition_lista.append(('\t'+tp+' part'+str(n)+' = ' + str(start) + '-' + str(start+l-1) + ';'))

	start += l 
	
	n+=1


l = 0
sl = len(super_matrix_dic)


with open(outfile,'w') as fila:

	for i in super_matrix_dic:

		seq = super_matrix_dic[i]

		name = i
		
		fila.write('>'+name+'\n')
		
		fila.write(seq+'\n')





with open(outfile_partition,'w') as fila:

	fila.write('#nexus\n')
	
	fila.write('begin sets;\n')
	
	for i in partition_lista:

		fila.write(i+'\n')
	
	fila.write('end;')
